# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging
import pickle

import matplotlib
matplotlib.use('Agg')

import numpy as np
import pandas as pd
import statsmodels.api as sm
from absl import app, flags

from coin.base.datetime_util import to_datetime

import experimental.prophet as prophet
from experimental.prophet.craps.universe import universe_all
from coin.experimental.jaewon.research.adaptive_edge import lt_model
from coin.experimental.jaewon.research.adaptive_edge import features
from coin.experimental.jaewon.research.adaptive_edge import y_features
from coin.experimental.jaewon.research.adaptive_edge.util import open_file, setup_plt, plot_bucket


def matrix(symbol_nick,
           ref_symbol_nicks,
           lt_model_param,
           order_latency,
           order_lifetime,
           taker_panalty_bp,
           date):
  vars = []

  # Timers, sampling points
  ts = prophet.global_timestamp()
  timer_1s = prophet.timer('1s')
  trigger = (ts % (4 * 10**9) == 0)
  vars += [ts]

  # Symbols
  universe = {elem.nickname: elem for elem in universe_all(date)}
  symbol_info = universe[symbol_nick]
  target_coin = symbol_info.product.base.currency.lower()
  assert target_coin in ('btc', 'eth', 'eos')

  feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)

  # Basic book info, mid-price Y
  vars += y_features.basic(symbol_info, feed)

  # Features
  with prophet.namespace('feature'):
    vars += features.features(target_coin, ref_symbol_nicks + [symbol_nick], date)

  # Sampler
  with prophet.namespace('sampler'):
    samplers = features.samplers(target_coin, 'feature/')
    vars += samplers

  # Trigger
  vars += [features.trigger(target_coin, symbol_nick, 'feature/')]

  # Long-term model theo price
  with prophet.namespace('lt_model'):
    theo_ask0_p, theo_bid0_p = lt_model.theo_price(symbol_info, date, **lt_model_param)
    vars += [theo_ask0_p, theo_bid0_p]

  # Fill info
  with prophet.namespace('fillinfo'):
    vars += y_features.multiple_edge_fillinfo(symbol_info,
                                              feed,
                                              theo_ask0_p,
                                              theo_bid0_p,
                                              trigger,
                                              order_latency,
                                              order_lifetime,
                                              taker_panalty_bp / 10000.)

  with prophet.control_if(trigger):
    aggregator = prophet.scalar_value_aggregator(vars)

  return aggregator


def make_y_inplace(df, y_windows=None):
  y_windows = y_windows or ['1s', '2s', '5s', '10s', '15s', '20s', '30s', '1m', '2m', '5m']

  cols = list(reversed(df.columns))

  buy_fill_cols = [c for c in cols if c.startswith('fillinfo/buy_filled_')]
  buy_taker_cols = [c for c in cols if c.startswith('fillinfo/buy_taker_')]
  buy_fp_cols = [c for c in cols if c.startswith('fillinfo/buy_fill_price_')]
  buy_fee_cols = [c for c in cols if c.startswith('fillinfo/buy_fee_')]
  buy_edges = [int(c[20:]) for c in buy_fill_cols]

  sell_fill_cols = [c for c in cols if c.startswith('fillinfo/sell_filled_')]
  sell_taker_cols = [c for c in cols if c.startswith('fillinfo/sell_taker_')]
  sell_fp_cols = [c for c in cols if c.startswith('fillinfo/sell_fill_price_')]
  sell_fee_cols = [c for c in cols if c.startswith('fillinfo/sell_fee_')]
  sell_edges = [int(c[21:]) for c in sell_fill_cols]

  def _int(x):
    return int(x) if isinstance(x, str) else x

  for w in y_windows:
    logging.info('Gen Y: %s' % w)
    y_mid = df['Y/mid_p.Y%s' % w]

    # Buy
    buy_ret_df = pd.DataFrame(index=df.index)
    for idx, c in enumerate(buy_fp_cols):
      fp_col = c
      fee_col = buy_fee_cols[idx]
      colname = '%d' % buy_edges[idx]
      buy_ret_df[colname] = (y_mid - df[fp_col] - df[fee_col]) * df[buy_fill_cols[idx]].replace(
          0., np.nan)
    df['Y/buy_edge_Y%s' % w] = [_int(v) for v in buy_ret_df.idxmax(axis=1)]
    df['Y/buy_ret_Y%s' % w] = buy_ret_df.max(axis=1)

    neg_ret_mask = (df['Y/buy_ret_Y%s' % w] < 0.)
    df['Y/buy_edge_Y%s' % w][neg_ret_mask] += 1.

    # Sell
    sell_ret_df = pd.DataFrame(index=df.index)
    for idx, c in enumerate(sell_fp_cols):
      fp_col = c
      fee_col = sell_fee_cols[idx]
      colname = '%d' % sell_edges[idx]
      sell_ret_df[colname] = (df[fp_col] - y_mid - df[fee_col]) * df[sell_fill_cols[idx]].replace(
          0., np.nan)
    df['Y/sell_edge_Y%s' % w] = [_int(v) for v in sell_ret_df.idxmax(axis=1)]
    df['Y/sell_ret_Y%s' % w] = sell_ret_df.max(axis=1)

    neg_ret_mask = (df['Y/sell_ret_Y%s' % w] < 0.)
    df['Y/sell_edge_Y%s' % w][neg_ret_mask] += 1.


def train_single_sampler(df, x_cols, y_col):
  X = df[x_cols]
  Y = df[y_col]

  mod = sm.GLS(Y, X)
  res = mod.fit()
  summary = res.summary()
  res.remove_data()

  return {'model': res, 'summary': summary}


def train(df, out_path, sampler_cols, is_ratio=0.7):
  x_cols = [c for c in df.columns if c.startswith('feature/') and not c.endswith('.x')] + ['const']
  y_cols = [c for c in df.columns if c.startswith('Y/buy_edge_Y') or c.startswith('Y/sell_edge_Y')]

  # Remove na.
  all_cols = sampler_cols + x_cols + y_cols + ['global_timestamp']
  XY = sm.add_constant(df)[all_cols]
  logging.info('XY.shape = %s', str(XY.shape))
  mask = ~XY.isna().any(axis=1)
  XY = XY[mask]
  logging.info('XY[mask].shape = %s', str(XY.shape))

  XY_is = XY.iloc[:int(XY.shape[0] * is_ratio)]
  XY_os = XY.iloc[int(XY.shape[0] * is_ratio):]
  logging.info('XY_is.shape=%s XY_os.shape=%s', XY_is.shape, XY_os.shape)
  if XY_os.shape[0] > 0:
    print(type(XY_os['global_timestamp'].iloc[0]))
    logging.info('Out-sample starting: %s', to_datetime(int(XY_os['global_timestamp'].iloc[0])))

  models = {}
  for scol in sampler_cols:
    sampler_mask = (XY_is[scol] == 1.)
    buy_model = train_single_sampler(XY_is[sampler_mask], x_cols=x_cols, y_col='Y/buy_edge_Y20s')
    sell_model = train_single_sampler(XY_is[sampler_mask], x_cols=x_cols, y_col='Y/sell_edge_Y20s')

    models[scol] = {'buy': buy_model, 'sell': sell_model}

  with open_file(out_path, 'wb') as f_out:
    pickle.dump(models, f_out)
  '''
  ys = {
      'buy_Y_is': XY_is['Y/buy_edge_Y15s'],
      'buy_Y_is_pred': buy_model['model'].predict(XY_is[x_cols]),
      'buy_Y_os': XY_os['Y/buy_edge_Y15s'],
      'buy_Y_os_pred': buy_model['model'].predict(XY_os[x_cols]),

      'sell_Y_is': XY_is['Y/sell_edge_Y15s'],
      'sell_Y_is_pred': sell_model['model'].predict(XY_is[x_cols]),
      'sell_Y_os': XY_os['Y/sell_edge_Y15s'],
      'sell_Y_os_pred': sell_model['model'].predict(XY_os[x_cols]),
  }
  '''
  return models


def main(argv):
  FLAGS = flags.FLAGS
  assert FLAGS.machine

  symbol_nick = FLAGS.symbol_nickname

  ref_symbol_nicks = []
  if FLAGS.machine == 'feed-01.ap-northeast-1.aws' and 'btc' in symbol_nick:
    ref_symbol_nicks.append('bitflyer-fxbtcjpy')

  logging.info('%s %s',
               FLAGS.matrix_out.format(nickname=symbol_nick),
               FLAGS.model_out.format(nickname=symbol_nick))

  lt_model_param = {'version': FLAGS.lt_model_version}
  order_latency = FLAGS.order_latency
  order_lifetime = FLAGS.order_lifetime

  end_date = datetime.datetime.now().date() - datetime.timedelta(days=2)
  begin_date = end_date - datetime.timedelta(days=30)
  logging.info('%s-%s', str(begin_date), str(end_date))

  friday_filter = lambda d: (d.weekday() != 4)
  taker_panalty_bp = 2.
  df = prophet.run_from_fastfeed_multiprocess(functools.partial(matrix,
                                                                symbol_nick,
                                                                ref_symbol_nicks,
                                                                lt_model_param,
                                                                order_latency,
                                                                order_lifetime,
                                                                taker_panalty_bp),
                                              begin_date,
                                              end_date,
                                              machine=FLAGS.machine,
                                              max_workers=None,
                                              debug_executor=FLAGS.debug_executor,
                                              inject_date_to_graph_func=True,
                                              date_filter=friday_filter,
                                              duration_after='20m')

  make_y_inplace(df)

  print(df.describe())

  if FLAGS.matrix_out:
    df.to_hdf(FLAGS.matrix_out.format(nickname=symbol_nick), '_')

  sampler_cols = [str(c) for c in df.columns if c.startswith('sampler/')]

  df['const'] = 1.
  model = train(df, FLAGS.model_out.format(nickname=symbol_nick), sampler_cols, 0.99)


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_string(
      'symbol_nickname',
      'okex-btcusd-perpetual',
      'Symbol nickname such as '
      'huobi-btcusd-quarter, okex-btcusd-quarter, okex-btcusd-perpetual, bitflyer-fxbtcjpy')

  flags.DEFINE_integer('lt_model_version', 1, 'LT Model version')

  flags.DEFINE_string('order_latency', '0.5s', 'Order latency')

  flags.DEFINE_string('order_lifetime', '30s', 'Order life')

  flags.DEFINE_boolean('debug_executor', False, 'Debug executor')

  flags.DEFINE_string('matrix_out', 'dump_{nickname}.hd5.lz4', 'Matrix output filepath')

  flags.DEFINE_string('model_out', 'model_{nickname}.pkl.gz', 'Model output filepath')

  flags.DEFINE_string('machine',
                      None,
                      'Machine such as feed-01.ap-northeast-1.aws, feed-01.cn-hongkong.aliyun')

  app.run(main)
